{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numpy.fft import fft\n",
    "from numpy import linspace, sin, pi, power, ceil, log2, arange, random\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fft_func(fs, data):\n",
    "    len_ = len(data)\n",
    "    n = int(power(2, ceil(log2(len_))))\n",
    "    fft_y_ = (fft(y, n)) / len_ * 2\n",
    "    fre_ = arange(int(n / 2)) * fs / n\n",
    "    fft_y_ = fft_y_[range(int(n / 2))]\n",
    "    return fre_, fft_y_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == '__main__':\n",
    "    Fs = 256000\n",
    "    f1 = 390\n",
    "    f2 = 2e3\n",
    "\n",
    "    noise1 = random.random(Fs)\n",
    "    noise2 = random.normal(1, 10, Fs)\n",
    "\n",
    "    t = linspace(0, 1, Fs)\n",
    "    y = 2 * sin(2 * pi * f1 * t) + 5 * sin(2 * pi * f2 * t)\n",
    "    # y = 2 * sin(2 * pi * f1 * t) + 5 * sin(2 * pi * f2 * t) + noise2\n",
    "\n",
    "    fre, fft_y = fft_func(Fs, y)\n",
    "    plt.figure()\n",
    "    plt.plot(fre, abs(fft_y))\n",
    "    plt.grid()\n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
